// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//


#pragma once

#include <melon/singleton.h>
#include <memory>

/// StatsReporter designed to assist in reporting various metrics of the
/// application that uses pollux library. The library itself does not implement
/// the StatsReporter and it should be implemented by the application.
///
/// To initialize the reporter singleton in your application use this pattern
/// (note that MyReporter should implement the abstract class
/// BaseStatsReporter):
///
///   melon::Singleton<kumo::pollux::BaseStatsReporter> reporter([]() {
///     return new MyReporter();
///   });
///
/// Then, for every metric that needs to be reported, it is required to register
/// one (usually) or more types (StatType) before reporting the metric:
///
///   DEFINE_METRIC("my_stat1", kumo::pollux::StatType::COUNT);
///
/// To register one histogram, it requires the min and max value of
///  the range, the bucket width as well as the percentiles to be reported.
///   DEFINE_HISTOGRAM_METRIC("my_stat2", 10, 0, 100, 50, 99, 100);
///
/// The StatType controls how metric is aggregated.
/// After that, every call to RECORD_METRIC_VALUE increases the metric by the
/// given value:
///
///   By default the following will add 1 to the metric if not provided value
///   RECORD_METRIC_VALUE("my_stat1");
///   RECORD_METRIC_VALUE("my_stat2", 10);
///   RECORD_METRIC_VALUE("my_stat1", numOfFailures);

namespace kumo::pollux {
    enum class StatType {
        /// Tracks the average of the inserted values.
        AVG,
        /// Tracks the sum of the inserted values.
        SUM,
        /// Tracks the sum of the inserted values per second.
        RATE,
        /// Tracks the count of inserted values.
        COUNT,
        /// Tracks the histogram of inserted values.
        HISTOGRAM,
    };

    inline std::string statTypeString(StatType stat) {
        switch (stat) {
            case StatType::AVG:
                return "Avg";
            case StatType::SUM:
                return "Sum";
            case StatType::RATE:
                return "Rate";
            case StatType::COUNT:
                return "Count";
            case StatType::HISTOGRAM:
                return "Histogram";
            default:
                return fmt::format("UNKNOWN: {}", static_cast<int>(stat));
        }
    }

    /// This is the base stats reporter interface that should be extended by
    /// different implementations.
    class BaseStatsReporter {
    public:
        virtual ~BaseStatsReporter() {
        }

        /// Register a stat of the given stat type.
        /// @param key The key to identify the stat.
        /// @param statType How the stat is aggregated.
        virtual void registerMetricExportType(const char *key, StatType statType)
        const = 0;

        virtual void registerMetricExportType(
            melon::StringPiece key,
            StatType statType) const = 0;

        /// Register a histogram with a list of percentiles defined.
        /// @param key The key to identify the histogram.
        /// @param bucketWidth The width of the buckets.
        /// @param min The starting value of the buckets.
        /// @param max The ending value of the buckets.
        /// @param pcts The aggregated percentiles to be reported.
        virtual void registerHistogramMetricExportType(
            const char *key,
            int64_t bucketWidth,
            int64_t min,
            int64_t max,
            const std::vector<int32_t> &pcts) const = 0;

        virtual void registerHistogramMetricExportType(
            melon::StringPiece key,
            int64_t bucketWidth,
            int64_t min,
            int64_t max,
            const std::vector<int32_t> &pcts) const = 0;

        /// Add the given value to the stat.
        virtual void addMetricValue(const std::string &key, size_t value = 1)
        const = 0;

        virtual void addMetricValue(const char *key, size_t value = 1) const = 0;

        virtual void addMetricValue(melon::StringPiece key, size_t value = 1)
        const = 0;

        /// Add the given value to the histogram.
        virtual void addHistogramMetricValue(const std::string &key, size_t value)
        const = 0;

        virtual void addHistogramMetricValue(const char *key, size_t value) const = 0;

        virtual void addHistogramMetricValue(melon::StringPiece key, size_t value)
        const = 0;

        /// Return the aggregated metrics in a serialized string format.
        virtual std::string fetchMetrics() = 0;

        static bool registered;
    };

    // This is a dummy reporter that does nothing
    class DummyStatsReporter : public BaseStatsReporter {
    public:
        void registerMetricExportType(const char * /*key*/, StatType /*statType*/)
        const override {
        }

        void registerMetricExportType(
            melon::StringPiece /*key*/,
            StatType /*statType*/) const override {
        }

        void registerHistogramMetricExportType(
            const char * /*key*/,
            int64_t /* bucketWidth */,
            int64_t /* min */,
            int64_t /* max */,
            const std::vector<int32_t> & /* pcts */) const override {
        }

        void registerHistogramMetricExportType(
            melon::StringPiece /* key */,
            int64_t /* bucketWidth */,
            int64_t /* min */,
            int64_t /* max */,
            const std::vector<int32_t> & /* pcts */) const override {
        }

        void addMetricValue(const std::string & /* key */, size_t /* value */)
        const override {
        }

        void addMetricValue(const char * /* key */, size_t /* value */)
        const override {
        }

        void addMetricValue(melon::StringPiece /* key */, size_t /* value */)
        const override {
        }

        void addHistogramMetricValue(const std::string & /* key */, size_t /* value */)
        const override {
        }

        void addHistogramMetricValue(const char * /* key */, size_t /* value */)
        const override {
        }

        void addHistogramMetricValue(melon::StringPiece /* key */, size_t /* value */)
        const override {
        }

        std::string fetchMetrics() override {
            return "";
        }
    };

#define DEFINE_METRIC(key, type)                               \
  {                                                            \
    if (::kumo::pollux::BaseStatsReporter::registered) {    \
      auto reporter = melon::Singleton<                        \
          kumo::pollux::BaseStatsReporter>::try_get_fast(); \
      if (MELON_LIKELY(reporter != nullptr)) {                 \
        reporter->registerMetricExportType((key), (type));     \
      }                                                        \
    }                                                          \
  }

#define RECORD_METRIC_VALUE(key, ...)                          \
  {                                                            \
    if (::kumo::pollux::BaseStatsReporter::registered) {    \
      auto reporter = melon::Singleton<                        \
          kumo::pollux::BaseStatsReporter>::try_get_fast(); \
      if (MELON_LIKELY(reporter != nullptr)) {                 \
        reporter->addMetricValue((key), ##__VA_ARGS__);        \
      }                                                        \
    }                                                          \
  }

#define DEFINE_HISTOGRAM_METRIC(key, bucket, min, max, ...)    \
  {                                                            \
    if (::kumo::pollux::BaseStatsReporter::registered) {    \
      auto reporter = melon::Singleton<                        \
          kumo::pollux::BaseStatsReporter>::try_get_fast(); \
      if (MELON_LIKELY(reporter != nullptr)) {                 \
        reporter->registerHistogramMetricExportType(           \
            (key),                                             \
            (bucket),                                          \
            (min),                                             \
            (max),                                             \
            (std::vector<int32_t>({__VA_ARGS__})));            \
      }                                                        \
    }                                                          \
  }

#define RECORD_HISTOGRAM_METRIC_VALUE(key, ...)                  \
  {                                                              \
    if (::kumo::pollux::BaseStatsReporter::registered) {      \
      auto reporter = melon::Singleton<                          \
          kumo::pollux::BaseStatsReporter>::try_get_fast();   \
      if (MELON_LIKELY(reporter != nullptr)) {                   \
        reporter->addHistogramMetricValue((key), ##__VA_ARGS__); \
      }                                                          \
    }                                                            \
  }
} // namespace kumo::pollux
